Preserve HF PTQ checkpoint sidecar files [NV BUG 6491822] - #2060

🔒 Security & Privacy | 🟠 Major | 🏗️ Heavy lift

#!/bin/bash
set -euo pipefail

tmp_dir="$(mktemp -d)"
trap 'rm -rf "$tmp_dir"' EXIT

mkdir -p "$tmp_dir/source" "$tmp_dir/export"
printf 'sensitive-data\n' > "$tmp_dir/secret.txt"
ln -s "$tmp_dir/secret.txt" "$tmp_dir/source/tokenizer.json"

SOURCE_DIR="$tmp_dir/source" EXPORT_DIR="$tmp_dir/export" python - <<'PY'
import os
import shutil
from pathlib import Path

source = Path(os.environ["SOURCE_DIR"]) / "tokenizer.json"
destination = Path(os.environ["EXPORT_DIR"]) / "tokenizer.json"

assert source.is_file(), "Path.is_file() follows the symlink"
shutil.copy2(source, destination)
assert destination.read_text() == "sensitive-data\n"
print("Symlink dereference reproduced.")
PY

# Inspect the changed copier and its tests before defining allowed symlink targets.
rg -n -C 4 '_should_copy_checkpoint_sidecar|copy2|is_file|copy_custom_model_files' \
  examples/hf_ptq/example_utils.py tests/examples/hf_ptq/test_example_utils.py

#!/bin/bash
set -euo pipefail

echo '--- helper and copier ---'
sed -n '900,1030p' examples/hf_ptq/example_utils.py

echo '--- relevant tests ---'
sed -n '1,180p' tests/examples/hf_ptq/test_example_utils.py

echo '--- symlink and cache references ---'
rg -n -C 3 'symlink|resolve\(|snapshot|cache|copy_custom_model_files|checkpoint sidecar' \
  examples tests SECURITY.md 2>/dev/null || true

#!/bin/bash
set -euo pipefail
log=/tmp/coderabbit-shell-logs/shell-output-Gb9q3G

echo '--- captured source section ---'
line=$(rg -n '^examples/hf_ptq/example_utils.py:965:' "$log" | head -1 | cut -d: -f1)
if [ -n "${line:-}" ]; then
  start=$((line-70)); end=$((line+75))
  sed -n "${start},${end}p" "$log"
fi

echo '--- captured test section ---'
line=$(rg -n '^tests/examples/hf_ptq/test_example_utils.py:52:' "$log" | head -1 | cut -d: -f1)
if [ -n "${line:-}" ]; then
  start=$((line-5)); end=$((line+75))
  sed -n "${start},${end}p" "$log"
fi

echo '--- focused repository search ---'
rg -n -C 3 'HF_HUB_CACHE|snapshot_download|local_files_only|is_symlink|readlink|resolve\(\)' \
  examples/hf_ptq tests/examples/hf_ptq SECURITY.md || true

#!/bin/bash
set -euo pipefail

echo '--- source ---'
sed -n '940,1030p' examples/hf_ptq/example_utils.py

echo '--- tests ---'
sed -n '1,125p' tests/examples/hf_ptq/test_example_utils.py

echo '--- focused references ---'
rg -n -C 3 'HF_HUB_CACHE|snapshot_download|local_files_only|is_symlink|readlink|resolve\(\)' \
  examples/hf_ptq tests/examples/hf_ptq SECURITY.md || true

#!/bin/bash
set -euo pipefail

echo '--- model-path resolution and cache roots ---'
sed -n '830,945p' examples/hf_ptq/example_utils.py

echo '--- security guidance ---'
rg -n -C 4 'checkpoint|sidecar|artifact|symlink|untrusted|deserialization' SECURITY.md

echo '--- call sites and export-path inputs ---'
rg -n -C 4 'copy_custom_model_files\(' examples/hf_ptq tests/examples/hf_ptq

Validate checkpoint sidecar symlink targets before copying.

Path.is_file() and shutil.copy2() follow symlinks. A local untrusted checkpoint can link a sidecar to any readable file outside the checkpoint and copy it into export_path. Reject targets outside approved checkpoint roots, while allowing Hugging Face cache snapshot-to-blob symlinks. Add an out-of-tree symlink regression test.

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@examples/hf_ptq/example_utils.py` around lines 965 - 970, Update
_should_copy_checkpoint_sidecar and the associated copy flow to validate symlink
targets before shutil.copy2: resolve the target and reject links escaping the
approved checkpoint roots, while permitting Hugging Face cache snapshot-to-blob
symlinks. Preserve existing exclusions for export-owned files and checkpoint
weight/index files, and add a regression test covering an out-of-tree sidecar
symlink.

Source: Path instructions

Bot comment.

*.gguf is missing from the weight-file pattern list (as are *.npz, *.pkl, *.mlmodel, *.tar*/*.zip archives). Many HF repos ship GGUF conversions alongside safetensors; with the blanket copy these get copied into the export dir, and via _resolve_model_path they'd also be downloaded (see comment on the snapshot_download call). Please extend the list.

Bot comment.

hf_ptq.py explicitly supports already-quantized sources (pack-quantized/compressed-tensors and MXFP4 gpt-oss). Those checkpoints often carry stale quant metadata sidecars — recipe.yaml (llm-compressor), quantize_config.json (AutoAWQ/GPTQ), quant_config.json — which the new blanket copy will drop next to the freshly written hf_quant_config.json, potentially confusing vLLM/SGLang scheme detection. Consider adding those to _EXPORT_OWNED_CHECKPOINT_FILES (or a separate "stale quant metadata" skip set).

Bot comment.

This changes an intentionally tiny fetch (allow_patterns=["*.py", "*.json"]) into "download every file in the repo that isn't a recognized weight file". For HF-ID sources that ship GGUF conversions, images/videos, or archives, this can be many GB downloaded at export time (and *.gguf isn't even in the ignore list). Suggest keeping an allow-list here (or adding a size/pattern guard) and broadening only the local copy step, which is what the bug is actually about.

fixed

Bot comment.

This reimplements an existing repo helper: modelopt/torch/export/plugins/hf_checkpoint_utils.py::copy_non_safetensor_files_from_ckpt already does "copy every non-safetensors file from a local HF checkpoint dir verbatim" and is the Megatron-Core path this PR says it wants to match (its docstring even documents which files modelopt owns). Since examples/hf_ptq already imports from modelopt.torch.export, consider calling it (extended with an exclude/extra-pattern argument for hf_quant_config.json + the broader weight patterns) instead of maintaining a second, subtly different copy routine.

Bot comment.

The old docstring deliberately gated executable code (modeling*.py, configuration_*.py, tokenization_*.py) behind trust_remote_code; after this change those files are copied into the export unconditionally and trust_remote_code only affects HF-ID resolution. That's likely intentional given the goal, but please call it out in the PR body (it's a behavior change relevant to SECURITY.md guidance) and add a test pinning the new parity between trust_remote_code=True/False.

Bot comment.

Good happy-path test. Missing cases worth adding: (a) subdirectories in the source (e.g. original/) are silently skipped — pin that so the loss of nested sidecars is a deliberate decision; (b) trust_remote_code=False now copies modeling_*.py, which is the actual behavior change; (c) the _resolve_model_path ignore-pattern change (mock snapshot_download and assert what is/ isn't requested), since that's the riskiest part of the diff.

🔒 Security & Privacy | 🟠 Major | 🏗️ Heavy lift

#!/bin/bash
set -eu

sed -n '280,330p' modelopt/torch/export/plugins/hf_checkpoint_utils.py
printf '\nRelevant callers:\n'
rg -n -A12 -B4 'copy_non_safetensor_files_from_ckpt|copy_custom_model_files' \
  modelopt examples tests

printf '\nPython standard-library behavior for the exact operations:\n'
python3 - <<'PY'
import os
import shutil
import tempfile
from pathlib import Path

with tempfile.TemporaryDirectory() as td:
    root = Path(td)
    source = root / "source"
    destination = root / "destination"
    source.mkdir()
    destination.mkdir()

    secret = root / "host-readable-secret.txt"
    secret.write_text("sensitive\n")
    sidecar = source / "tokenizer_config.json"
    sidecar.symlink_to(secret)

    print("isfile_symlink:", os.path.isfile(sidecar))
    shutil.copy2(sidecar, destination)
    copied = destination / sidecar.name
    print("copied_exists:", copied.exists())
    print("copied_is_symlink:", copied.is_symlink())
    print("copied_contents:", copied.read_text())
PY

Sensitive Data Exposure (CWE-59)

Reachability: External

● Entry
  examples/hf_ptq/example_utils.py:961
  copy_custom_model_files: Copy source checkpoint sidecar files to an HF PTQ export. The HF PTQ script writes ModelOpt-owned metadata and quantized weights first, t…
│
▼
● Sink
  modelopt/torch/export/plugins/hf_checkpoint_utils.py

Reject source symlinks before copying checkpoint sidecars.

os.path.isfile(sp) and shutil.copy2(sp, dst) follow symlinks. A sidecar link can copy a process-readable host file into the export. Reject symlinks or restrict targets to trusted snapshot roots, and add a regression test.

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@modelopt/torch/export/plugins/hf_checkpoint_utils.py` around lines 319 - 325,
Update the sidecar-copy loop around the os.path.isfile and shutil.copy2 calls to
reject symbolic-link sources before copying, using os.path.islink(sp) (or an
equivalent trusted-root validation) so only regular files from the snapshot are
exported. Preserve the existing exclusions and copy behavior for valid
non-symlink sidecars, and add a regression test covering a symlinked sidecar.

Source: Path instructions

[IMPORTANT Compatibility] The HF-ID path still uses an extension allowlist, so it reintroduces exactly the brittleness this PR removes from the local-dir path — and the two paths now disagree about what a "sidecar" is.

For a local --pyt_ckpt_path, copy_custom_model_files copies everything that isn't a weight/index/export-owned file. For an HF model ID, the set is bounded by what snapshot_download pulled, i.e. only *.json, *.md, *.model, *.py, *.tiktoken, *.txt. Files a real hub repo ships that never reach the export:

• chat_template.jinja — transformers ≥4.51 saves chat templates as a standalone .jinja sidecar, and many current hub repos ship one. hf_ptq.py:930-933 states the explicit intent that source tokenizer files win over regenerated ones "which may differ in format due to newer transformers versions". For .jinja that intent silently does not hold on the HF-ID path: the regenerated template is kept instead.
• LICENSE / NOTICE — extensionless, so never downloaded, even though the shared helper's own docstring names LICENSE as a file it exists to preserve.
• README assets (*.png, *.svg). Note the new test asserts accuracy_chart.png is copied — that assertion can only ever hold for a local source dir, so the test reads as broader coverage than it provides.

Why it matters: the reported bug (dropped reasoning parsers) is fixed for *.py, but the same class of bug remains for any sidecar whose extension isn't enumerated here — and it's invisible, since copy_custom_model_files reports success over an already-truncated directory.

Suggested fix: keep the allowlist only as a size guard and make it complete for non-weight text sidecars, e.g. add "*.jinja", "LICENSE*", "NOTICE*". Alternatively, invert to ignore_patterns=HF_CHECKPOINT_WEIGHT_FILE_PATTERNS on the download too, so both paths derive from the single weight-pattern source of truth (that list already covers *.gguf/*.npz/*.tar*/*.zip, which was the over-download concern from the earlier review round). Either way, worth a test that pins the download patterns and the copy exclusions against each other so they can't drift.

[IMPORTANT Compatibility] generation_config.json is excluded unconditionally, but only one of the two export paths that call copy_custom_model_files actually writes one — so on the TRT-LLM path this silently drops a file that used to be copied.

Tracing both call sites in hf_ptq.py:

• Unified HF path (hf_ptq.py:935): export_hf_checkpoint → model.save_pretrained writes generation_config.json. Excluding the source copy is correct here.
• TRT-LLM path (hf_ptq.py:884): export_tensorrt_llm_checkpoint writes only config.json and (for auto-quant) quant_cfg.json — grep for generation_config in modelopt/torch/export/model_config_export.py and tensorrt_llm_utils.py returns nothing. Before this PR, generation_config.json was copied here via the "*.json" code pattern whenever --trust_remote_code was set. Now it is never copied and never written, so the export loses the source's sampling defaults (temperature, top_p, eos_token_id overrides) outright.

Compare unified_export_megatron.py:319-330, which calls the same helper without exclude_files and then explicitly re-saves GenerationConfig.from_pretrained(...) — i.e. that path guarantees the file exists before excluding the source copy. The TRT-LLM path here does neither.

Same reasoning applies to config.json on the TRT-LLM path, but there the exclusion is right: the TRT-LLM config.json schema is unrelated to the HF one and copying the source over it would corrupt the export. generation_config.json has no such conflict.

Suggested fix: don't hardcode one exclusion set for both paths. Either pass the exclusions in from the caller so the TRT-LLM branch omits generation_config.json, or drop it from _HF_PTQ_EXPORT_OWNED_FILES and rely on ordering (the unified HF path runs save_pretrained before the copy, so the export version already exists — but note that means source-wins if you copy after, which is the opposite of today's behavior; pick one deliberately and say so in the docstring).

While here: the docstring at line 967-971 claims "Source processor files intentionally still win" — with preprocessor_config.json no longer in the exclusion set that is true, but the shared helper's docstring (hf_checkpoint_utils.py:298-299) lists preprocessor_config.json as modelopt-owned. Those two statements now contradict each other for readers of the shared helper.

[IMPORTANT Compatibility] This tuple is a new public name (no leading underscore) in a module that is star-exported through modelopt/torch/export/plugins/__init__.py (with import_plugin("hf_checkpoint_utils"): from .hf_checkpoint_utils import *), but the module has no __all__.

Two consequences:

1. Per CONTRIBUTING.md ("Define the public API with __all__ and re-export via from .module import *"), a new public symbol should be declared in __all__. Without one, the star-import also re-exports every module-level import — fnmatch, json, os, shutil, warnings, Iterable, Path, torch, snapshot_download, tqdm — into modelopt.torch.export.plugins. That's pre-existing, but this PR is the first to add an intentionally-public constant here, which makes the missing __all__ load-bearing rather than cosmetic.
2. Once exported it's an API surface with a compatibility obligation. If it's meant only for examples/hf_ptq (the sole consumer), name it _HF_CHECKPOINT_WEIGHT_FILE_PATTERNS or put it in __all__ deliberately.

Suggested fix: add __all__ = ["HF_CHECKPOINT_WEIGHT_FILE_PATTERNS", "copy_non_safetensor_files_from_ckpt", ...] naming the module's existing public functions, or rename the constant private if it isn't intended as public API.

Separately, on the list contents: *.npy/*.npz/*.pkl/*.pickle/*.zip/*.tar* aren't HF weight formats in the usual sense, and the name HF_CHECKPOINT_WEIGHT_FILE_PATTERNS oversells what it matches. Some legitimate sidecars are .npy (e.g. precomputed statistics shipped alongside a checkpoint) and would be silently dropped. Not blocking, but a name like ..._LARGE_ARTIFACT_PATTERNS would describe the intent (skip bulk binaries) more honestly than "weight files".

[SUGGESTION] Two small things in the copy loop:

1. The hardcoded safetensors check at line 321-322 is now redundant on the new path but inconsistent on the old one. HF_CHECKPOINT_WEIGHT_FILE_PATTERNS already covers *.safetensors and *.safetensors.index.json (which is a superset of the model.safetensors.index.json equality check — it also catches e.g. adapter_model.safetensors.index.json). For the Megatron caller, which passes no exclude_patterns, line 321 remains the only guard and still misses sharded index files not named exactly model.safetensors.index.json. Consider making HF_CHECKPOINT_WEIGHT_FILE_PATTERNS the always-applied baseline and letting exclude_patterns be purely additive — that removes the double source of truth and fixes the Megatron gap in one move.

2. Silent failures. The old copy_custom_model_files wrapped each shutil.copy2 in try/except and warned per file. The new loop does not, so a single unreadable file (bad permissions, broken symlink, dangling HF-cache blob) now aborts the whole sidecar copy — and because this runs after the weights are written, the user gets a traceback on an otherwise-complete export. A per-file try/except with a warning would preserve the previous resilience.

Both are non-blocking, but (2) is a real behavior regression for the HF-cache source case, where dangling blob symlinks do occur.

[SUGGESTION] This test doesn't exercise what its name claims. copy_custom_model_files passes trust_remote_code only to _resolve_model_path, which returns immediately at its first branch (if os.path.isdir(model_name_or_path): return model_name_or_path) for a local source_dir. So the flag is never read, and the assertion is identical to what the previous test already covers for super_v3_reasoning_parser.py with trust_remote_code=False.

That matters more than a redundancy nit here, because the trust_remote_code gate on executable *.py copying is precisely the behavior this PR changes: modeling*.py / configuration_*.py / tokenization_*.py are now copied into the export even for native (non-remote-code) loads, whereas the removed docstring explicitly justified gating them. A test named after trust_remote_code that can't observe the flag gives false confidence that the parity was checked.

Suggested fix: either drop this test (the first one covers .py copying), or make it actually meaningful — parametrize trust_remote_code=[True, False] over a modeling_custom.py fixture and assert the (now intentional) equality, so the behavior change is pinned as deliberate rather than incidental.

[SUGGESTION] Unlike the unified-HF call site at line 934-935, this TRT-LLM call is not guarded by if args.dist_state.is_main:. With world_size > 1, every rank races to shutil.copy2 the same sidecars into the same export_path, which can produce truncated files (a reader hitting a partially-written destination) as well as duplicated log spam.

CONTRIBUTING.md's "Develop with distributed processing in mind" calls this out explicitly: "Guard shared side effects, such as file writes or shared state updates, against race conditions between ranks."

The pre-existing code had the same gap, so this isn't introduced here — but this PR both broadens what gets copied (more files, larger files → wider race window) and touches this exact line, so it's a cheap fix to fold in:

            # Copy source checkpoint sidecar files for TensorRT-LLM export.
            if args.dist_state.is_main:
                copy_custom_model_files(args.pyt_ckpt_path, export_path, args.trust_remote_code)

Bot comment.

Unresolved from the previous round: generation_config.json is excluded for both callers, but only the unified-HF path writes one (save_pretrained). On the TRT-LLM path, export_tensorrt_llm_checkpoint writes only config.json (and quant_cfg.json for auto-quant) — there is no generation_config write anywhere in model_config_export.py. Before this PR the file was copied here via the old "*.json" pattern under --trust_remote_code; now it is neither copied nor written, so the export silently loses the source sampling defaults.

Compare unified_export_megatron.py, which calls the same helper with no exclude_files and then explicitly re-saves GenerationConfig.from_pretrained(...) — i.e. it guarantees the file exists before superseding the source copy.

Suggest either passing the exclusion set in from the caller (so the TRT-LLM branch keeps generation_config.json) or writing a fresh generation_config.json on that path. Excluding config.json on the TRT-LLM path is correct (different schema), so this needs to be per-caller rather than one shared set.

Bot comment.

This test can't observe the flag it's named after: copy_custom_model_files only forwards trust_remote_code to _resolve_model_path, which returns immediately for a local directory. So it's an exact duplicate of the .py-copy assertion in the preceding test.

That matters because the trust_remote_code gate on executable *.py is precisely what this PR changes — modeling*.py/configuration_*.py/tokenization_*.py are now copied for native (non-remote-code) loads, whereas the removed docstring explicitly justified gating them. Either drop this test, or parametrize trust_remote_code=[True, False] over a modeling_custom.py fixture and assert the (now intentional) equality so the change is pinned as deliberate. Please also call the behavior change out in the PR body — it's a security-posture change relative to SECURITY.md guidance.

Bot comment.

The old copy_custom_model_files wrapped each shutil.copy2 in try/except and warned per file. The new loop doesn't, so a single unreadable entry (bad permissions, dangling HF-cache blob symlink) aborts the whole sidecar copy — and since this runs after the quantized weights are written, the user gets a traceback on an otherwise-complete export. A per-file try/except + warning would preserve the previous resilience.

Also, the hardcoded entry.endswith(".safetensors") or entry == "model.safetensors.index.json" check is now a second source of truth: for the Megatron caller (no exclude_patterns) it still misses index files not named exactly model.safetensors.index.json. Consider making HF_CHECKPOINT_WEIGHT_FILE_PATTERNS the always-applied baseline and keeping exclude_patterns purely additive.

Bot comment.

HF_CHECKPOINT_WEIGHT_FILE_PATTERNS is a new public name in a module star-exported via plugins/__init__.py, and this module has no __all__ (so the star-import also re-exports fnmatch/json/os/shutil/torch/...). Per CONTRIBUTING.md, declare the public API with __all__, or make the constant private (_HF_...) if examples/hf_ptq is the only intended consumer — otherwise it becomes an API surface with a compat obligation.

Minor naming nit: *.npy/*.npz/*.pkl/*.zip/*.tar* aren't HF weight formats; some legitimate sidecars are .npy (precomputed stats) and will now be silently dropped. Something like ..._LARGE_ARTIFACT_PATTERNS would describe the intent more honestly.

Bot comment.

Critical: removing the hardcoded entry.endswith(".safetensors") or entry == "model.safetensors.index.json" guard leaves this helper with no default exclusions, and the weight-pattern list now lives in examples/hf_ptq rather than here.

The other caller, unified_export_megatron.py::GPTModelExporter.save_pretrained, invokes copy_non_safetensor_files_from_ckpt(self._hf_pretrained_model_name, save_directory) with no exclude_files/exclude_patterns. With this change it will copy the source BF16 *.safetensors shards and model.safetensors.index.json into save_directory immediately before save_safetensors_by_layer_index writes the quantized shards + index. Consequences: (a) the full unquantized checkpoint is duplicated into the export (disk/time), and (b) if the source shard count differs from the exported one, stale model-000XX-of-000YY.safetensors files and/or a stale index survive alongside the new ones → a checkpoint that loads wrong or fails to load.

Suggested fix (also suggested in an earlier round): make a module-level weight/bulk-artifact pattern tuple the always-applied baseline inside this function and let exclude_patterns be purely additive. That restores Megatron safety, closes the pre-existing *.safetensors.index.json naming gap, and lets examples/hf_ptq pass only its extra patterns. The docstring ("Copy every non-safetensors file…") and the function name are also inaccurate as written.

Bot comment.

This assertion should fail against the current implementation: after this PR the default (no-exclusion) call copies everything, including model.safetensors — exclude_files is an empty set and exclude_patterns an empty tuple, and the hardcoded .safetensors skip was deleted. Please confirm the unit suite was run.

The assertion itself is the right contract; it's the implementation that drifted (see comment on hf_checkpoint_utils.py). Once the weight patterns are applied as a library-side baseline, this test passes and also pins the Megatron caller's behavior. While here, consider asserting pytorch_model.bin/stats.npy presence/absence deliberately: the current expectations (copied by default, excluded only when the example passes patterns) are a behavior decision worth documenting in the test name/comment.